SageMaker Python SDKのS3モジュールを使って、S3スキームのURLでのS3オブジェクト操作を楽にする

SageMaker Python SDKを用いることでSageMakerを使った機械学習モデルの開発を効率化できます。 SageMakerを使った機械学習モデルの開発をする上で、S3とのデータのやりとりは避けて通れません。しかも、SageMaker Python SDKでは、Boto3のS3クライアントやリソースでは直接扱えない、S3スキーマのURL形式が使われています。

今回はそういったS3に関連した操作を楽にしてくれる、SageMaker Python SDKに含まれるS3のユーティリティモジュールについて紹介します。

• S3 Utilities — sagemaker 1.50.6 documentation

S3 Utilities

S3のユーティリティモジュールはsagemaker.s3として定義されており、その中には次のような関数とクラスが定義されています。

• parse_s3_url
  • S3スキームのURLをバケットとキーに変換します。
  • 例: 's3://bucket/path/to/object' -> ('bucket', 'path/to/object')
• S3Uploader
  • ローカルファイルもしくは文字列として保持しているデータをS3にアップロードする静的メソッドが定義されたクラスです。
  • 各メソッドではS3スキームのURLでアップロード先を指定します。
• S3Downloader
  • S3上のオブジェクト一覧を取得したり、S3上のオブジェクトをダウンロードする静的メソッドが定義されたクラスです。
  • 各メソッドではS3スキームのURLでS3上の場所を指定します。

sagemaker.s3の内容はシンプルで限定的ですが、各関数やメソッドにはS3スキームのURLを使用するという共通要素があります。Boto3ではバケット名とキーでS3上のオブジェクトを指定しますが、SageMaker Python SDKではS3スキームのURL形式が用いられています。そのため、SageMaker Python SDKを使ってモデルの開発をする際には、"S3スキームのURL"と"バケット名とキー"それぞれの形式を相互に変換する必要が出てきます。 sagemaker.s3を利用することで、S3スキームのURLそのままでオブジェクトのS3へのアップロードやS3からのダウンロードが行えます。

やってみる

parse_s3_urlを使って、S3スキームのURLからバケットとキーに変換します。

import sagemaker

sagemaker.s3.parse_s3_url('s3://bucket/path/to/object.txt')

S3Uploader.upload_string_as_file_bodyを使って、CSV形式の文字列をアップロードします。

import pandas as pd
df = pd.DataFrame([[1,2,3], [2,4,5]], columns=['a', 'b', 'c'])

csv_s3_uri = sagemaker.s3.S3Uploader.upload_string_as_file_body(
    body=df.to_csv(index=False),
    desired_s3_uri='s3://bucket/upload-test/x.csv'
)
sagemaker.s3.S3Downloader.list(csv_s3_uri)

S3Uploader.uploadを使って、複数のCSVファイルをアップロードします。

!mkdir test_data
for i in range(10):
    df.to_csv(f'test_data/{i}.csv', index=False)
csv_dir_s3_uri = 's3://bucket/test_data'
sagemaker.s3.S3Uploader.upload(
    local_path='./test_data/',
    desired_s3_uri=csv_dir_s3_uri
)
sagemaker.s3.S3Downloader.list(csv_dir_s3_uri)

S3Downloader.downloadで先ほどアップロードしたCSVファイルをダウンロードしてみます。

local_path = './downloaded_data/'
sagemaker.s3.S3Downloader.download(
    s3_uri=csv_dir_s3_uri,
    local_path=local_path
)
!ls $local_path

さいごに

SageMaker Python SDKに含まれるS3のユーティリティモジュールについて紹介しました。"S3スキームのURL"と"バケット名とキー"の変換は簡単な処理で実現できますが、何回も繰り返しているとかなり面倒になります。sagemaker.s3モジュールを利用することで、そういった面倒を緩和することができそうです。

参考

• sagemaker-python-sdk/s3.py at master · aws/sagemaker-python-sdk